package com.util;

import lombok.extern.log4j.Log4j2;

import javax.websocket.Session;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author ruiwen
 * @date 2024/11/13 11:40
 */
@Log4j2
public class WebsocketUtil {

    private static Map<String, Session> ONLINE_SESSION = new ConcurrentHashMap<>();

    /**
     * 添加session
     */
    public static void addSession(String pcUniqueId, Session session) {
        // 一个用户只允许一个session链接
        ONLINE_SESSION.putIfAbsent(pcUniqueId, session);
        log.info("pcUniqueId [{}] connected. Total online pcUniqueIds: {}", pcUniqueId, ONLINE_SESSION.size());
    }

    /**
     * 移除session
     */
    public static void removeSession(String pcUniqueId) {
        ONLINE_SESSION.remove(pcUniqueId);
        log.info("PC [{}] disconnected. Total online PCs: {}", pcUniqueId, ONLINE_SESSION.size());
    }

    /**
     * 给单个用户推送消息
     */
    public static void sendMessage(String pcUniqueId, String message) {
        Session session = ONLINE_SESSION.get(pcUniqueId);
        if (session == null) {
            log.warn("Session for PC [{}] not found", pcUniqueId);
            return;
        }
        sendMessage(session, message);
    }

    public static void sendMessage(Session session, String message) {
        if (session != null) {
            session.getAsyncRemote().sendText(message);
        }
    }

    /**
     * 给所有用户发消息
     */
    public static void sendMessageForAll(String message) {
        ONLINE_SESSION.forEach((pcUniqueId, session) -> {
            CompletableFuture.runAsync(() -> sendMessage(session, message))
                    .exceptionally(ex -> {
                        log.error("Failed to send message to PC [{}]: {}", pcUniqueId, ex.getMessage());
                        return null;
                    });
        });
    }

    /**
     * 给指定的多个用户推送消息
     */
    public static void sendMessageForPCs(Set<String> pcUniqueIds, String message) {
        pcUniqueIds.forEach(pcUniqueId -> {
            Session session = ONLINE_SESSION.get(pcUniqueId);
            if (session == null) {
                log.warn("Session for PC [{}] not found", pcUniqueId);
                return;
            }
            CompletableFuture.runAsync(() -> sendMessage(session, message))
                    .exceptionally(ex -> {
                        log.error("Failed to send message to PC [{}]: {}", pcUniqueId, ex.getMessage());
                        return null;
                    });
        });
    }

    public static List<String> getAllKeys() {
        //把ONLINE_SESSION key 转换成list
        return new ArrayList<>(ONLINE_SESSION.keySet());

    }


}
